{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "EgiF12Hf1Dhs"
   },
   "source": [
    "This notebook provides examples to go along with the [textbook](http://manipulation.csail.mit.edu/trajectories.html).  I recommend having both windows open, side-by-side!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "eeMrMI0-1Dhu"
   },
   "outputs": [],
   "source": [
    "from functools import partial\n",
    "\n",
    "import numpy as np\n",
    "from pydrake.all import (\n",
    "    AddMultibodyPlantSceneGraph,\n",
    "    DiagramBuilder,\n",
    "    InverseKinematics,\n",
    "    MeshcatVisualizer,\n",
    "    MeshcatVisualizerParams,\n",
    "    Parser,\n",
    "    Rgba,\n",
    "    RigidTransform,\n",
    "    Solve,\n",
    "    Sphere,\n",
    "    StartMeshcat,\n",
    ")\n",
    "\n",
    "from manipulation import running_as_notebook\n",
    "from manipulation.scenarios import AddPlanarIiwa, AddShape, AddWsg\n",
    "from manipulation.utils import ConfigureParser"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Start the visualizer.\n",
    "meshcat = StartMeshcat()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Visualizing the configuration space\n",
    "\n",
    "I've got the default sampling resolution set fairly fine; it will take a minute or so to run."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def draw_configuration_space(shelves=True):\n",
    "    builder = DiagramBuilder()\n",
    "\n",
    "    plant, scene_graph = AddMultibodyPlantSceneGraph(builder, time_step=0.001)\n",
    "    iiwa = AddPlanarIiwa(plant)\n",
    "    wsg = AddWsg(plant, iiwa, roll=0.0, welded=True)\n",
    "    sphere = AddShape(plant, Sphere(0.02), \"sphere\")\n",
    "    X_WO = RigidTransform([0.6, 0, 0.65])\n",
    "    plant.WeldFrames(plant.world_frame(), plant.GetFrameByName(\"sphere\"), X_WO)\n",
    "\n",
    "    if shelves:\n",
    "        parser = Parser(plant)\n",
    "        ConfigureParser(parser)\n",
    "        bin = parser.AddModelsFromUrl(\"package://manipulation/shelves.sdf\")[0]\n",
    "        plant.WeldFrames(\n",
    "            plant.world_frame(),\n",
    "            plant.GetFrameByName(\"shelves_body\", bin),\n",
    "            RigidTransform([0.6, 0, 0.4]),\n",
    "        )\n",
    "\n",
    "    plant.Finalize()\n",
    "\n",
    "    #    visualizer = MeshcatVisualizer.AddToBuilder(\n",
    "    #        builder,\n",
    "    #        scene_graph,\n",
    "    #        meshcat,\n",
    "    #        MeshcatVisualizerParams(delete_prefix_initialization_event=False))\n",
    "\n",
    "    diagram = builder.Build()\n",
    "    context = diagram.CreateDefaultContext()\n",
    "    plant_context = plant.GetMyContextFromRoot(context)\n",
    "\n",
    "    q0 = plant.GetPositions(plant_context)\n",
    "    gripper_frame = plant.GetFrameByName(\"body\", wsg)\n",
    "    sphere_frame = plant.GetFrameByName(\"sphere\", sphere)\n",
    "\n",
    "    ik = InverseKinematics(plant, plant_context)\n",
    "    collision_constraint = ik.AddMinimumDistanceLowerBoundConstraint(\n",
    "        0.001, 0.01\n",
    "    )\n",
    "    grasp_constraint = ik.AddPositionConstraint(\n",
    "        gripper_frame, [0, 0.1, 0], sphere_frame, [0, 0, 0], [0, 0, 0]\n",
    "    )\n",
    "\n",
    "    prog = ik.get_mutable_prog()\n",
    "    q = ik.q()\n",
    "    prog.SetInitialGuess(q, q0)\n",
    "    result = Solve(ik.prog())\n",
    "    if not result.is_success():\n",
    "        print(\"IK failed\")\n",
    "\n",
    "    diagram.ForcedPublish(context)\n",
    "\n",
    "    def eval(q0, q1, q2, c, tol):\n",
    "        return float(c.evaluator().CheckSatisfied([q0, q1, q2], tol))\n",
    "\n",
    "    meshcat.Delete()\n",
    "    meshcat.SetProperty(\"/Background\", \"visible\", False)\n",
    "    meshcat.SetObject(\"initial_guess\", Sphere(0.05), Rgba(0.4, 1, 1, 1))\n",
    "    meshcat.SetTransform(\"initial_guess\", RigidTransform(q0))\n",
    "    meshcat.SetObject(\"ik_solution\", Sphere(0.05), Rgba(0.4, 0.4, 1, 1))\n",
    "    meshcat.SetTransform(\"ik_solution\", RigidTransform(result.GetSolution(q)))\n",
    "\n",
    "    low = plant.GetPositionLowerLimits()\n",
    "    up = plant.GetPositionUpperLimits()\n",
    "    N = 70 if running_as_notebook else 5\n",
    "    vertices, triangles = mcubes.marching_cubes_func(\n",
    "        tuple(low),\n",
    "        tuple(up),\n",
    "        N,\n",
    "        N,\n",
    "        N,\n",
    "        partial(eval, c=grasp_constraint, tol=0.05),\n",
    "        0.5,\n",
    "    )\n",
    "    meshcat.SetTriangleMesh(\n",
    "        \"grasp_constraint\", vertices.T, triangles.T, Rgba(0.5, 0.9, 0.5)\n",
    "    )\n",
    "\n",
    "    if shelves:\n",
    "        vertices, triangles = mcubes.marching_cubes_func(\n",
    "            tuple(low),\n",
    "            tuple(up),\n",
    "            N,\n",
    "            N,\n",
    "            N,\n",
    "            partial(eval, c=collision_constraint, tol=0.0),\n",
    "            0.5,\n",
    "        )\n",
    "        meshcat.SetTriangleMesh(\n",
    "            \"collision_constraint\",\n",
    "            vertices.T,\n",
    "            triangles.T,\n",
    "            Rgba(0.9, 0.5, 0.5, 1),\n",
    "        )\n",
    "\n",
    "\n",
    "# I've made pymcubes optional (it's a heavy dependency)\n",
    "try:\n",
    "    import mcubes\n",
    "\n",
    "    draw_configuration_space(shelves=True)\n",
    "except ImportError:\n",
    "    pass"
   ]
  }
 ],
 "metadata": {
  "colab": {
   "collapsed_sections": [],
   "name": "Robotic Manipulation - Motion Planning.ipynb",
   "provenance": []
  },
  "kernelspec": {
   "display_name": "Python 3.8.10 64-bit",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.5"
  },
  "vscode": {
   "interpreter": {
    "hash": "e7370f93d1d0cde622a1f8e1c04877d8463912d04d973331ad4851f04de6915a"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}
